# -*- coding: utf-8 -*-
"""
 Copyright (c) [2021] Huawei Technologies Co.,Ltd.ALL rights reserved.
 This program is licensed under Mulan PSL v2.
 You can use it according to the terms and conditions of the Mulan PSL v2.
          http://license.coscl.org.cn/MulanPSL2
 THIS PROGRAM IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 See the Mulan PSL v2 for more details.

 @Author  : lemon-higgins
 @email   : lemon.higgins@aliyun.com
 @Date    : 2021-04-22 10:52:19
 @License : Mulan PSL v2
 @Version : 1.0
 @Desc    : 测试设备名获取
"""

import os
import sys
import subprocess
import argparse

SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
sys.path.append(SCRIPT_PATH)
import ssh_cmd
import mugen_log
import rpm_manage


def get_test_nic(node=1):
    """获取可测试使用的网卡

    Args:
        node (int, optional): 节点号. Defaults to 1.

    Returns:
        [str]: 网卡名
    """
    exitcode, tmpfile = rpm_manage.rpm_install(pkgs="lshw", node=node)
    if exitcode != 0:
        mugen_log.logging(
            "error",
            "Failed to install the dependent software package required to obtain the network card.",
        )
        sys.exit(exitcode)

    if os.environ.get("NODE" + str(node) + "_LOCALTION") == "local":
        dev_info = subprocess.getoutput("cat /proc/device-tree/model")
        if "Lichee Pi 4A" in dev_info:
            # 兼容  licheepi4a 硬件板
            output = subprocess.getoutput(
                "ip addr show | grep '^[0-9]:'|grep -v lo |grep -v virbr |grep -v wlan0 |grep -v dummy0 |grep -v sit0@NONE |grep -v ip_vti0@NONE |cut -d ':' -f2 |sed -e 's/^[[:space:]]*//' | grep -v '"
                + os.environ.get("NODE" + str(node) + "_NIC")
                + "'"
            ).replace("\n", " ")
        else:
            output = subprocess.getoutput(
                "ip link | grep 'state UP' | awk -F': ' '{print $2}' | grep -v '"
                + os.environ.get("NODE" + str(node) + "_NIC")
                + "'"
            ).replace("\n", " ")
    else:
        conn = ssh_cmd.pssh_conn(
            os.environ.get("NODE" + str(node) + "_IPV4"),
            os.environ.get("NODE" + str(node) + "_PASSWORD"),
            os.environ.get("NODE" + str(node) + "_SSH_PORT"),
            os.environ.get("NODE" + str(node) + "_USER"),
        )
        dev_info = ssh_cmd.pssh_cmd(
            conn,
            "cat /proc/device-tree/model"
        )[1]
        if "Lichee Pi 4A" in dev_info:
            # 兼容 licheepi4a 硬件板
            output = ssh_cmd.pssh_cmd(
                conn,
                "ip addr show | grep '^[0-9]:'|grep -v lo |grep -v virbr |grep -v wlan0 |grep -v dummy0 |grep -v sit0@NONE |grep -v ip_vti0@NONE |cut -d ':' -f2 |sed -e 's/^[[:space:]]*//' | grep -v '"
                + os.environ.get("NODE" + str(node) + "_NIC")
                + "'",
            )[1].replace("\n", " ")
        else:
            output = ssh_cmd.pssh_cmd(
                conn,
                "ip link | grep 'state UP' | awk -F': ' '{print $2}' | grep -v '"
                + os.environ.get("NODE" + str(node) + "_NIC")
                + "'",
            )[1].replace("\n", " ")

        ssh_cmd.pssh_close(conn)

    if tmpfile is not None:
        rpm_manage.rpm_remove(node=node, tmpfile=tmpfile)

    return output


def get_test_disk(node=1):
    """获取可测试使用的磁盘

    Args:
        node (int, optional): 节点号. Defaults to 1.

    Returns:
        [str]: 磁盘名称
    """
    if os.environ.get("NODE" + str(node) + "_LOCALTION") == "local":
        used_disk = subprocess.getoutput(
            r"lsblk -l | grep -e '/.*\|\[.*\]' | awk '{print $1}' | tr -d '[0-9]' | uniq | sed -e ':a;N;$!ba;s/\\n/ /g'"
        )

        test_disk = subprocess.getoutput(
            r"lsblk -n | grep 'disk' | grep -v 'mmcblk.*\|mtdblock.*\|nbd.*\|└─.*\|"
            + used_disk.replace(" ", "\|")
            + r"' | awk '{print $1}' | sed -e ':a;N;$!ba;s/\\n/ /g'"
        )
    else:
        conn = ssh_cmd.pssh_conn(
            os.environ.get("NODE" + str(node) + "_IPV4"),
            os.environ.get("NODE" + str(node) + "_PASSWORD"),
            os.environ.get("NODE" + str(node) + "_SSH_PORT"),
            os.environ.get("NODE" + str(node) + "_USER"),
        )
        used_disk = ssh_cmd.pssh_cmd(
            conn,
            r"lsblk -l | grep -e '/.*\|\[.*\]' | awk '{print $1}' | tr -d '[0-9]' | uniq | sed -e ':a;N;$!ba;s/\\n/ /g'",
        )[1]
        test_disk = ssh_cmd.pssh_cmd(
            conn,
            r"lsblk -n | grep 'disk' | grep -v 'mmcblk.*\|mtdblock.*\|nbd.*\|└─.*\|"
            + used_disk.replace(" ", "\|")
            + r"' | awk '{print $1}' | sed -e ':a;N;$!ba;s/\\n/ /g'",
        )[1]

    return test_disk


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="manual to this script")
    parser.add_argument("--node", type=int, default=1)
    parser.add_argument("--device", type=str, choices=["nic", "disk"], default="nic")
    args = parser.parse_args()

    if args.device == "nic":
        print(get_test_nic(args.node))
    elif args.device == "disk":
        print(get_test_disk(args.node))
    else:
        mugen_log.logging(
            "warn",
            "No other test driven acquisition is provided at this time, you can issue to us for follow-up.",
        )
        sys.exit(1)
